package com.hyperf.core.lock.executor;

import com.alibaba.fastjson.JSON;
import org.apache.curator.shaded.com.google.common.collect.Maps;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author lizhaoyang
 */
public class RedisInterProcessMutex {

    private static final RedisScript<String> SCRIPT_UNLOCK = new DefaultRedisScript<>("if redis.call('get',KEYS[1]) " +
            "== ARGV[1] then return tostring(redis.call('del', KEYS[1])==1) else return 'false' end", String.class);

    private final StringRedisTemplate stringRedisTemplate;
    private final String lockKey;
    private final String lockValue;
    private final ConcurrentMap<Thread, LockData> threadData;

    public RedisInterProcessMutex(StringRedisTemplate stringRedisTemplate, String lockKey, String lockValue) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.lockKey = lockKey;
        this.lockValue = lockValue;
        this.threadData = Maps.newConcurrentMap();
    }

    public boolean acquire(long time, TimeUnit unit) {
        Thread currentThread = Thread.currentThread();
        LockData lockData = this.threadData.get(currentThread);
        if (lockData != null) {
            lockData.lockCount.incrementAndGet();
            return true;
        } else {
            Boolean lockResult = stringRedisTemplate.opsForValue().setIfAbsent(this.lockKey, this.lockValue, time, unit);
            if (lockResult == null || Boolean.FALSE.equals(lockResult)) {
                return false;
            } else {
                lockData = new LockData(currentThread, lockKey, lockValue);
                threadData.put(currentThread, lockData);
                return true;
            }
        }
    }

    public void release() throws Exception {
        Thread currentThread = Thread.currentThread();
        LockData lockData = this.threadData.get(currentThread);
        if (lockData == null) {
            throw new IllegalMonitorStateException("You do not own the lock: " + this.lockKey);
        } else {
            int newLockCount = lockData.lockCount.decrementAndGet();
            if (newLockCount <= 0) {
                if (newLockCount < 0) {
                    throw new IllegalMonitorStateException("Lock count has gone negative for lock: " + this.lockKey);
                } else {
                    try {
                        String removeResult = stringRedisTemplate.execute(SCRIPT_UNLOCK, Collections.singletonList(lockKey), lockValue);
                        if (!Boolean.parseBoolean(removeResult)) {
                            throw new IllegalMonitorStateException("delete lock error: " + this.lockKey);
                        }
                    } finally {
                        this.threadData.remove(currentThread);
                    }
                }
            }
        }
    }

    private static class LockData {
        private final Thread owningThread;
        private final String lockKey;
        private final String lockValue;
        private final AtomicInteger lockCount;

        private LockData(Thread owningThread, String lockKey, String lockValue) {
            this.owningThread = owningThread;
            this.lockKey = lockKey;
            this.lockValue = lockValue;
            this.lockCount = new AtomicInteger(1);
        }

        @Override
        public String toString() {
            return JSON.toJSONString(this);
        }
    }
}
